In [ ]:
#    This code computes the transformation rules to fortran code and makes other checks.
#
#    Copyright (c) 2015 by Malte Titze (malte.titze@cern.ch)
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

from sympy import *
import numpy as np
import math
import cmath # only used later for C. Iselins formulas
import textwrap as tw
from IPython.display import Math, display, Image
from __future__ import division   # to make sure that 1/2 = 0.5 just in case...

import copy

# To handle a rational p/q, use Rational(p,q)

# Define certain variables we need later
o1, o2, o3, o4, o5, o6 = symbols("o1 o2 o3 o4 o5 o6", real=True)
o = [o1, o2, o3, o4, o5, o6]
def orbit(i):
    return o[i-1]

x, px, z, pz = symbols("x p_x z p_z", real=True)
sig = Symbol("\sigma", real=True)
psig = Symbol("p_\sigma", real=True)
rkn = [x, px, z, pz, sig, psig]
def ripken(i):
    return rkn[i-1]

kx, kz, g, g2 = symbols("K_x K_z g g2", real=True)
sg = Symbol("sstr", real=True)
og = Symbol("oct", real=True)

dipi, dipr = symbols("dipi dipr", real=True)
beta0 = Symbol("beta0", positive=True)
elrad = Symbol("\Delta s", positive=True)


hx = Function("h_x")(x, z)
hz = Function("h_z")(x, z)
etah = Function("etahat", real=True)(psig)
kappa = Function("kappa", real=True)(x, px, z, pz, psig)

In [ ]:
def f_wrapper(str):
    # fortran string wrapper in case this is necessary
    
    increment = 8  # in case the output must be pasted in code with increment
    lines = tw.wrap(str, 72 - increment)
    output = ''
    for line in lines:
        output = output + line + '&\n     &'
    
    # remove the last empty line
    output = output[:-8]
    
    return output


def elrad_order(expr, order):
    # This function returns an expression which is 'expr' developed up to 'order' in elrad.
    # If the order is negative (e.g. -1), then it returns the expression.
    
    result = expr
    
    if order != oo:
        result = expr.subs(elrad, 0)
        
        expr2 = expr
        for i in range(order):
            expr2 = expr2.diff(elrad)
            result = result + Rational(1, math.factorial(i + 1))*expr2.subs(elrad, 0)*elrad**(i + 1)
        
    return simplify(result, ratio=1)


def madxconverter(expr):
    # madxconverter conversts expression 'expr' which is in ripkens notation (see literature) to 'orbit' notation
    
    result = 0
    if expr != 0:
        # step 1: express the derivative of eta hat by eta and psigma:
        result = expr.subs(etah.diff(psig),(1 + beta0**2*psig)/(1 + etah))
    
        # step 2: express all coordinates in terms of orbit
        result = result\
        .subs(x, o1)\
        .subs(px, o2)\
        .subs(z, o3)\
        .subs(pz, o4)\
        .subs(sig, o5*beta0)\
        .subs(psig, o6/beta0)
    
    return N(result)


def inv_madxconverter(expr):
    # inv_madxconverter converts orbit notation to ripken notation
    # warning: expr should not contain any derivative expressions!
    
    result = expr\
    .subs(o1, x)\
    .subs(o2, px)\
    .subs(o3, z)\
    .subs(o4, pz)\
    .subs(o5, sig/beta0)\
    .subs(o6, psig*beta0)
    
    return N(result)


def at_reference_point(expr, reference_point):
    '''
    This function substitutes a reference point into expr. It was created in order to evaluate
    the derivatives at 0 (which corresponds to the MADX-Implementation)
    '''
    eta_at_reference_point5 = 0   # careful: The full formula should be used if one sets psig != 0
    
    result = expr.subs(o1, reference_point[0]).subs(o2, reference_point[1])\
    .subs(o3, reference_point[2]).subs(o4, reference_point[3])\
    .subs(o5, reference_point[4]).subs(o6, reference_point[5])

    result = result.subs(etah.subs(psig, reference_point[5]), eta_at_reference_point5)
    return result


def f_multipole2(x2, px2, z2, pz2, sig2, psig2, bx_hor_plane, bz_hor_plane):
    # f_multipole2 computes the vector potential of the combined function magnet
    
    # bx_hor_plane, bz_hor_plane are the field distribution by which we determine the gauge of the potential.
    #
    # Examples:
    # bx_hor_plane = [-kz]
    # bx_hor_plane = [0]
    #
    # bz_hor_plane = [kx, g]
    # bz_hor_plane = [kx, g, Rational(1,2)*sg]
    # bz_hor_plane = [kx, g, Rational(1,2)*sg, Rational(1,6)*og]
    # bz_hor_plane = [0, g]
    
    # Get K_x and K_z:
    if len(bz_hor_plane) == 0:
        bz_hor_plane = [0]
        
    if len(bx_hor_plane) == 0:
        bx_hor_plane = [0]
        
    kx2 = bz_hor_plane[0]
    kz2 = -bx_hor_plane[0]
        
    # fastnirty True: skip the Maxwell-conformal condition and use only the field distribution conditions above
    fastndirty = False
    
    # compute add_ord additional orders - works only if fastndirty = False
    add_ord = 0
    
    # construct gamma-conditions vector, n0 and N (see theory)
    gamma = bz_hor_plane[:]
    j = 0
    n0 = 0
    for ele in bx_hor_plane:
        try:
            gamma[j] = gamma[j] + I*ele
        except IndexError:
            gamma.append(I*ele)
        if (gamma[j] == 0) and (n0 == j):
            n0 = n0 + 1
    
        j = j + 1
    
    N = len(gamma) - 1
    n0 = min(n0, N) # to handle the exeption that N = 0, n0 = 1
    # n0 is the lowest index which is not zero; N the highest which is not zero.
    
    # compute F
    rp = Rational(1,2)*(x2 + I*z2)
    rm = conjugate(rp)
    sigma = kx2 + I*kz2
    csigma = conjugate(sigma)
    
    F = 0
    gk = [[0]*(N + add_ord + 1)]
    for k in range(n0, N + 2 + add_ord):   # k = n0, ..., N + 1 + add_ord
        
        gkp1 = [0]*(k + 2)
        if (n0 == k):
            gkp1[0] = gamma[n0]
        if (n0 < k) and (k <= N):
            gkp1[0] = gamma[k] + kx2*gamma[k - 1]
        if (k == N + 1):
            gkp1[0] = kx2*gamma[N]
        gkp1[0] = gkp1[0]*Rational(-2**k, k + 1)
        
        GK1, j = 0, 0
        while j < k and not fastndirty:   # j = 0, ..., k - 1
            # note that gkp1[k-j] = conjugate(gkp1[j+1]), this can be used to 'optimize' this loop.
            gkp1[j+1] = csigma*gk[-1][j+1]*Rational(j - k + Rational(3,2),k - j) + \
            sigma*gk[-1][j]*Rational(-j + Rational(1,2), j + 1)

            gkp1[0] = gkp1[0] - gkp1[j+1]*Rational(k - j, k + 1)
            GK1 = GK1 + gkp1[j+1]*rp**(k - j)*rm**(j + 1)
            
            j = j + 1
        
        gkp1[k+1] = conjugate(gkp1[0])
        gk.append(gkp1)
        
        # GK1 = GK1 + 2*re(gkp1[0]*rp**(k + 1))
        GK1 = GK1 + gkp1[0]*rp**(k + 1) + conjugate(gkp1[0]*rp**(k + 1))  # 2*re(gkp1[0]*rp**(k+1)) may give buggy code (some type problems)
        
        F = F + GK1
    
    return expand(F)


def hamil(x, px, z, pz, sig, psig, bx_hor_plane, bz_hor_plane):
    # the Hamiltonian
    kx2 = bz_hor_plane[0]
    kz2 = -bx_hor_plane[0]
    
    result = psig - (1 + kx2*x + kz2*z)*((1 + etah)**2 - px**2 - pz**2)**Rational(1,2)\
    - f_multipole2(x, px, z, pz, sig, psig, bx_hor_plane, bz_hor_plane)
    return result


def S(x2, px2, z2, pz2, sig2, psig2, bx_hor_plane, bz_hor_plane):
    # S is the function h according to the definitions in my handout.
    gg_expr = f_multipole2(x, px, z, pz, sig, psig, bx_hor_plane, bz_hor_plane)
    ux = px + elrad*gg_expr.diff(x)
    uz = pz + elrad*gg_expr.diff(z)
    
    kx2 = bz_hor_plane[0]
    kz2 = -bx_hor_plane[0]
    
    factor = 1/(1 + elrad**2*(kx2**2 + kz2**2))
    xi = 2*elrad*(ux*kx2 + uz*kz2)*factor
    zeta = (ux**2 + uz**2 - (1 + etah)**2)*factor
    
    result = -Rational(1,2)*xi + sqrt(Rational(1,4)*xi**2 - zeta)
    #result = sqrt((1 + etah)**2 - px**2 - pz**2)    
    # note that the reference point must be in orbit-notation
    result = at_reference_point(madxconverter(result), [x2, px2, z2, pz2, sig2/beta0, psig2*beta0])
    
    return result


def S0(x2, px2, z2, pz2, sig2, psig2):
    # S0 is the function h_0 according to the definitions in my handout.
    
    result = sqrt((1 + etah)**2 - px2**2 - pz2**2)    
    # note that the reference point must be in orbit-notation
    result = at_reference_point(madxconverter(result), [x2, px2, z2, pz2, sig2/beta0, psig2*beta0])
    
    return result


def concat2(re, te, re2, te2):
    # this function computes the 2nd derivative of the composition of the maps belonging to (re, te)
    # and (re2, te2) respectively, i.e. it computes the operation (re, te) * (re2, te2)
    
    npre = np.array(re)
    npte = np.array(te)
    npre2 = np.array(re2)
    npte2 = np.array(te2)

    result_re = np.tensordot(npre, npre2, (1,0))
    result_te = np.tensordot(np.tensordot(npte, npre2, (1,0)), npre2, (1,0)) + np.tensordot(npre, npte2, (1,0))
    
    return [result_re, result_te]


def concat2b(re, te, re2, te2):
    # the same as concat2, but a more direct (slower) implementation
    npre = np.array(re)
    npte = np.array(te)
    npre2 = np.array(re2)
    npte2 = np.array(te2)
    
    result_re = [[0 for i in range(6)] for j in range(6)] # initialization
    for i in range(6):
        for j in range(6):
            sum1 = 0
            for k in range(6):
                sum1 = sum1 + npre[i][k]*npre2[k][j]
            result_re[i][j] = sum1
    
    result_te = [[[0 for i in range(6)] for j in range(6)] for k in range(6)]   # initialization    
    for i in range(6):
        for j in range(6):
            for k in range(6):
                sum1 = 0
                for a in range(6):
                    sum2 = 0
                    for b in range(6):
                        sum2 = sum2 + npte[i][a][b]*npre2[a][j]*npre2[b][k]
                    sum1 = sum1 + sum2 + npre[i][a]*npte2[a][j][k]
                result_te[i][j][k] = sum1
                
    return [result_re, result_te]

In [ ]:
# Define the field distribution by which we determine the gauge of the potential. Examples:

# Horizontal:
# -----------
# bx_hor_plane = [-kz]
bx_hor_plane = [0]

# Vertical:
# ---------
# bz_hor_plane = [0]
# bz_hor_plane = [kx]
# bz_hor_plane = [0, g]
bz_hor_plane = [kx, g]
# bz_hor_plane = [kx, g, Rational(1,2)*sg]
# bz_hor_plane = [kx, g, Rational(1,2)*sg, Rational(1,6)*og]  

if len(bz_hor_plane) == 0:
    bz_hor_plane = [0]
        
if len(bx_hor_plane) == 0:
    bx_hor_plane = [0]
        
kx2 = bz_hor_plane[0]
kz2 = -bx_hor_plane[0]

gg_expr = f_multipole2(x, px, z, pz, sig, psig, bx_hor_plane, bz_hor_plane)

#gg_expr2 = g*(-Rational(1,2)*x**2 - Rational(1,3)*kx*x**3 + Rational(1,2)*z**2)
#gg_expr_bend = -Rational(1,2)*kx**2*x**2 - kx*x
#gg_expr = gg_expr2 + gg_expr_bend

nshow = False  # True: show the Potential function (1 + kx*x + kz*z)*A_s and the Magnetic field in the midplane z == 0
if nshow:
    pure_latex = False
    mystr = '\\begin{align*}' + \
                 'G &= ' + latex( gg_expr ) + '\\\\' + \
                 'B_x\\mid_{z=0} &= ' + latex(expand((gg_expr\
                                              .diff(z)/(1 + kx2*x + kz2*z)).subs(z, 0)) ) + '\\\\' + \
                 'B_z\\mid_{z=0} &= ' + latex(-expand((gg_expr\
                                               .diff(x)/(1 + kx2*x + kz2*z)).subs(z, 0)) ) + '\\\\' + \
            '\\end{align*}'\
            
    if pure_latex:
        print mystr
    else:
        display(Math(mystr))

In [ ]:
def get_order(ham, order):
    # straightforward implementation of 1st- and 2nd order (non-symplectic) transformation rules
    # input:
    # order = 1 or 2
    # ham: hamiltonian
    
    myzero = 0*x  # trick to be able to differentiate integer 0
    
    dqh = [ham.diff(x), ham.diff(z), myzero]
    
    dpsigh = ham.diff(psig).subs(etah.diff(psig),(1 + beta0**2*psig)/(1 + etah))
    dph = simplify([ham.diff(px), ham.diff(pz), dpsigh])
    
    q = [x, z, sig]
    p = [px, pz, psig]
    
    dpph = [[dph[i].diff(p[k]).subs(etah.diff(psig),(1 + beta0**2*psig)/(1 + etah)) for k in range(3)] for i in range(3)]
    dqph = [[dph[i].diff(q[k]) for k in range(3)] for i in range(3)]
    
    dpqh = [[dqh[i].diff(p[k]).subs(etah.diff(psig),(1 + beta0**2*psig)/(1 + etah)) for k in range(3)] for i in range(3)]
    dqqh = [[dqh[i].diff(q[k]) for k in range(3)] for i in range(3)]
    
    # initialization
    dq, dp = 3*[0], 3*[0]
    
    for i in range(3):
        dq[i] = q[i] + elrad*dph[i]
        if order >= 2:
            for k in range(3):
                dq[i] = dq[i] - elrad**2*(Rational(3,2)*dpph[k][i]*dqh[k] + Rational(1,2)*dpqh[i][k]*dph[k])
            
    for i in range(3):
        dp[i] = p[i] - elrad*dqh[i]
        if order >= 2:
            for k in range(3):
                dp[i] = dp[i] + elrad**2*(Rational(3,2)*dpqh[k][i]*dqh[k] + Rational(1,2)*dpqh[i][k]*dph[k])
    
    return simplify([dq, dp])


def at_zero(expr):
    return expr.subs(kx, 0).subs(kz, 0).subs(g, 0).subs(g2, 0).subs(sg, 0).subs(og, 0)

def get_kd_order(ham, order):
    # straightforward implementation of 1st- and 2nd order drift-kick-drift transformation rules.
    # input:
    # order = 1 or 2
    # ham: hamiltonian
    
    myzero = 0*x  # trick to be able to differentiate integer 0
    
    dqh = [ham.diff(x), ham.diff(z), myzero]
    
    dpsigh = ham.diff(psig).subs(etah.diff(psig),(1 + beta0**2*psig)/(1 + etah))
    dph = simplify([ham.diff(px), ham.diff(pz), dpsigh])
    
    q = [x, z, sig]
    p = [px, pz, psig]
    
    dpph = [[dph[i].diff(p[k]).subs(etah.diff(psig),(1 + beta0**2*psig)/(1 + etah)) for k in range(3)] for i in range(3)]
    dqph = [[dph[i].diff(q[k]) for k in range(3)] for i in range(3)]
    
    dpqh = [[dqh[i].diff(p[k]).subs(etah.diff(psig),(1 + beta0**2*psig)/(1 + etah)) for k in range(3)] for i in range(3)]
    dqqh = [[dqh[i].diff(q[k]) for k in range(3)] for i in range(3)]
    
    # initialization
    dq, dp = 3*[0], 3*[0]
    
    for i in range(3):
        dq[i] = q[i] + elrad*dph[i] - elrad*at_zero(dph[i])
        if order >= 2:
            for k in range(3):
                dq[i] = dq[i] - elrad**2*(Rational(3,2)*dpph[k][i]*dqh[k] + Rational(1,2)*dpqh[i][k]*dph[k]) \
                + elrad**2*(dqph[i][k]*at_zero(dqh[k]) + dpph[i][k]*at_zero(dph[k]))
            
    for i in range(3):
        dp[i] = p[i] - elrad*dqh[i] + elrad*at_zero(dqh[i])
        if order >= 2:
            for k in range(3):
                dp[i] = dp[i] + elrad**2*(Rational(3,2)*dpqh[k][i]*dqh[k] + Rational(1,2)*dpqh[i][k]*dph[k]) \
                - elrad**2*(dqqh[i][k]*at_zero(dqh[k]) + dpqh[i][k]*at_zero(dph[k]))
    
    return simplify([dq, dp])

In [ ]:
nshow = False # True: show the Hamiltonian and the transformation rules

# Set option for transformation formula:
#
# 0: thick lens formula (required for iten-function to compare the values with madx output),
#   however, this option yield very complicated functions for the program to handle, it is therefore recommended to
#   use the wrong, but simplified option 4.
# 1: kick formula (for tracking)
# 2: plain drift (for testing purposes; alternatively set the fields to zero
# 3: first-order drift-kick-drift map (see my handout)
# 4: simplified version of the full formula of option 0.
# 5: second-order drift-kick-drift map (under development!)
# 6: Iselin drift-kick.
# 7: Iselin thick map
# 8: see 6
# 10: computed 1st and 2nd thick map (set the order there)
# 11: same as 10, but for kick-drift map. For example, option 11 with order = 1 corresponds to hand-computed
#     formula of option 3.

option = 10


if option == 0:
    toshow = 'thick lens formula'
    # compute the transformation
    u = px + elrad*gg_expr.diff(x)
    v = pz + elrad*gg_expr.diff(z)
    S_expr = S(x, px, z, pz, sig, psig, bx_hor_plane, bz_hor_plane)

    ripken_com = [\
              x + elrad*(1 + kx2*x + kz2*z)*(u/S_expr + elrad*kx2), \
              u + elrad*kx2*S_expr, \
              z + elrad*(1 + kx2*x + kz2*z)*(v/S_expr + elrad*kz2), \
              v + elrad*kz2*S_expr, \
              sig + elrad - elrad*(1 + kx2*x + kz2*z)*(1 + beta0**2*psig)/S_expr, \
              psig\
                 ]
    
    ripken_used = ripken_com
    

if option == 1:
    toshow = 'kick map'
    # compute the transformation
    u = px + elrad*gg_expr.diff(x)
    v = pz + elrad*gg_expr.diff(z)
    S_expr = S(x, px, z, pz, sig, psig, bx_hor_plane, bz_hor_plane)

    kick_map = [\
              x, \
              px + elrad*(kx2*(1 + etah) + gg_expr.diff(x)), \
              z, \
              pz + elrad*(kz2*(1 + etah) + gg_expr.diff(z)), \
              sig - elrad*(kx2*x + kz2*z)*(1 + beta0**2*psig)/(1 + etah), \
              psig\
                 ]
    
    ripken_used = kick_map

    
if option == 2:
    toshow = 'plain drift'
    
    x_drift = x + elrad*px*((1 + etah)**2 - px**2 - pz**2)**Rational(-1,2)
    px_drift = px
    z_drift = z + elrad*pz*((1 + etah)**2 - px**2 - pz**2)**Rational(-1,2)
    pz_drift = pz
    sig_drift = sig + elrad - elrad*(1 + beta0**2*psig)*((1 + etah)**2 - px**2 - pz**2)**Rational(-1,2)
    psig_drift = psig

    ripken_drift = [x_drift, px_drift, z_drift, pz_drift, sig_drift, psig_drift]
    
    ripken_used = ripken_drift
    
    
if option == 3:
    toshow = 'first-order drift-kick-drift map'
    # get first-order map
    u = px + elrad*gg_expr.diff(x)
    v = pz + elrad*gg_expr.diff(z)
    S0_expr = S0(x, px, z, pz, sig, psig)
    
    dkd1 = [\
              x + elrad*(kx2*x + kz2*z)*px/S0_expr, \
              u + elrad*kx2*S0_expr, \
              z + elrad*(kx2*x + kz2*z)*pz/S0_expr, \
              v + elrad*kz2*S0_expr, \
              sig + elrad - elrad*(kx2*x + kz2*z)*(1 + beta0**2*psig)/S0_expr, \
              psig\
                 ]
    
    ripken_used = dkd1
    

    
if option == 4:
    toshow = 'simplified thick-lens formula for iten'
    # this transformation is not correct, and only used for the purpose of providing a thick slice formula
    # to compute the concatenated re- and te-entries later, via the iten functions.
    u = px + elrad*gg_expr.diff(x)
    v = pz + elrad*gg_expr.diff(z)
    S0_expr = S0(x, px, z, pz, sig, psig)
    S_expr = S(x, px, z, pz, sig, psig, bx_hor_plane, bz_hor_plane)

    simple_formula = [\
              x + elrad*(1 + kx2*x + kz2*z)*(u/S0_expr + elrad*kx2), \
              u + elrad*kx2*S_expr, \
              z + elrad*(1 + kx2*x + kz2*z)*(v/S0_expr + elrad*kz2), \
              v + elrad*kz2*S_expr, \
              sig + elrad - elrad*(1 + kx2*x + kz2*z)*(1 + beta0**2*psig)/S0_expr, \
              psig\
                 ]
    
    ripken_used = simple_formula
    
    
if option == 5:
    toshow = 'second-order drift-kick-drift map'
    
    dgx = gg_expr.diff(x)
    dgz = gg_expr.diff(z)
    
    u = px + elrad*dgx
    v = pz + elrad*dgz
    S0_expr = S0(x, px, z, pz, sig, psig)
    
    # zero'th- and first order
    dkd2 = [\
            x + elrad*(kx2*x + kz2*z)*px/S0_expr, \
            u + elrad*kx2*S0_expr, \
            z + elrad*(kx2*x + kz2*z)*pz/S0_expr, \
            v + elrad*kz2*S0_expr, \
            sig + elrad - elrad*(kx2*x + kz2*z)*(1 + beta0**2*psig)/S0_expr, \
            psig\
           ]
    
    # add second-order
    alpha_x = kx2 + dgx/S0_expr
    alpha_z = kz2 + dgz/S0_expr
    
    dgxx = dgx.diff(x)
    dgxz = dgx.diff(z)
    dgzz = dgz.diff(z)
    
    dkd2[0] = dkd2[0] + elrad**2*Rational(1,2)*((1 + kx2*x + kz2*z)*\
    (alpha_x - alpha_z + px/S0_expr**2*(3*(px*alpha_x + pz*alpha_z) - px*kx2 - pz*kz2))\
                                               - kx2*(px/S0_expr)**2 - kz2*px*pz/S0_expr**2)
    
    dkd2[1] = dkd2[1] + elrad**2*Rational(1,2)*((px*alpha_x + pz*alpha_z)**2/S0_expr - kx2**2*S0_expr\
                                               - dgxx*px/S0_expr - dgxz*pz/S0_expr)
    
    dkd2[2] = dkd2[2] + elrad**2*Rational(1,2)*((1 + kx2*x + kz2*z)*\
    (alpha_z - alpha_x + pz/S0_expr**2*(3*(pz*alpha_z + px*alpha_x) - pz*kz2 - px*kx2))\
                                               - kz2*(pz/S0_expr)**2 - kx2*pz*px/S0_expr**2)
    
    dkd2[3] = dkd2[3] + elrad**2*Rational(1,2)*((pz*alpha_z + px*alpha_x)**2/S0_expr - kz2**2*S0_expr\
                                               - dgzz*pz/S0_expr - dgxz*px/S0_expr)
    
    dkd2[4] = dkd2[4] + elrad**2*Rational(1,2)*(1 + beta0**2*psig)*((1 + kx2*x + kz2*z)*(- 2*px*kx2 - 2*pz*kz2
        - 2*px*dgx/S0_expr - 2*pz*dgz/S0_expr + (px*dgx + pz*dgz)/S0_expr**3)\
                                                                    + (kx2*px + kz2*pz)/S0_expr**2)

    ripken_used = dkd2
    
    
if option not in [1, 2, 3, 4, 5]:
    # thick lens comparison to Iselin (1985) = thick result in MAD-X
    # first step: we define the first derivative
    # ripken_used = [0]*6
    # kx2 = curvature term h in his paper

    kax = sqrt(kx2**2 + g)
    kay = sqrt(-g)

    cx = cosh(1j*kax*elrad)
    cy = cosh(1j*kay*elrad)

    sx = sinh(1j*kax*elrad)/(1j*kax)
    sy = sinh(1j*kay*elrad)/(1j*kay)

    dx = (1.0 - cx)/kax**2

    j1 = (elrad - sx)/kax**2
    

    gamma = 1/sqrt(1 - beta0**2)

    m_iselin = Matrix([
            [cx, sx, 0, 0, 0, kx2/beta0*dx],
            [-kax**2*sx, cx, 0, 0, 0, kx2/beta0*sx],
            [0, 0, cy, sy, 0, 0],
            [0, 0, -kay**2*sy, cy, 0, 0],
            [-kx2/beta0*sx, -kx2/beta0*dx, 0, 0, 1, -(kx2/beta0)**2*j1 + elrad/(beta0*gamma)**2],
            [0, 0, 0, 0, 0, 1]
        ])
    # the matrix m_iselin is symplectic by construction (also checked with script below)
    
    # Next, we combine inverse drift operations from left and right with respect to elrad/2.
    # The derivative of the inverse of a plain drift, di_drift, evaluated at zero (Iselins formulas hold at zero):
    
    # elrad2 = Rational(1,2)*elrad   # !!!! drift-kick-drift
    elrad2 = elrad  # However, MAD-X alternates only drift and kick.
    
    di_drift = Matrix([
            [1, -elrad2, 0, 0, 0, 0],
            [0, 1, 0, 0, 0, 0],
            [0, 0, 1, -elrad2, 0, 0],
            [0, 0, 0, 1, 0, 0],
            [0, 0, 0, 0, 1, -elrad2/gamma**2],
            [0, 0, 0, 0, 0, 1]
        ])
    
    # kick_map = di_drift*m_iselin*di_drift
    kick_map = m_iselin*di_drift   # only drift-kick composition
    
    # The above formulas are verified in MAD-X
    ###################################################
    # now the second-order entries. According to Iselin
    
    t_iselin = [[[0 for i in range(6)] for j in range(6)] for k in range(6)]   # initialization
    
    gk1 = g
    gk2 = sg  # K1 iselin = g; K2 iselin = sg or sg/2; this is something I need to check.
    # kx2 = curvature term h in his paper
    
    j2 = Rational(1,2)*(3*elrad - 4*sx + sx*cx)/kax**4
    jd = ((1 - cosh(1j*2*kay*elrad))/(2*kay)**2 - (1 - cosh(1j*kax*elrad))/kax**2)/(kax**2 - 4*kay**2)
    js = (sinh(1j*2*kay*elrad)/(1j*2*kay) - sinh(1j*kax*elrad)/(1j*kax))/(kax**2 - 4*kay**2)
    jc = (cosh(1j*2*kay*elrad) - cosh(1j*kax*elrad))/(kax**2 - 4*kay**2)
    jf = ((elrad - sinh(1j*2*kay*elrad)/(1j*2*kay))/(1j*2*kay)**2 - \
         (elrad - sinh(1j*kax*elrad)/(1j*kax))/kax**2)/(kax**2 - 4*kay**2)
    j3 = (15*elrad - 22*sx + 9*sx*cx - 2*sx*cx**2)/(6*kx**6)
    
    k2pk1 = gk2 + 2*kx2*gk1
    
    # I interpret his notation '(h**2/4beta0)...' as division by 4*beta0
    t_iselin[0][0][0] = -Rational(1,6)*k2pk1*(sx**2 + dx) - Rational(1,2)*kx2*kax**2*sx**2
    t_iselin[0][0][1] = -Rational(1,6)*k2pk1*sx*dx + Rational(1,2)*kx2*sx*cx
    t_iselin[0][1][1] = -Rational(1,6)*k2pk1*dx**2 + Rational(1,2)*kx2*cx*dx
    t_iselin[0][0][5] = -Rational(1,12)*kx2/beta0*k2pk1*(3*sx*j1 - dx**2) + \
        Rational(1,2)*kx2**2/beta0*sx**2 + Rational(1,4)/beta0*gk1*elrad*sx
    t_iselin[0][1][5] = -Rational(1,12)*kx2/beta0*k2pk1*(sx*dx**2 - 2*cx*j2) + \
        Rational(1,4)*kx2**2/beta0*(sx*dx + cx*j1) - Rational(1,4)/beta0*(sx + elrad*cx)
    t_iselin[0][5][5] = -Rational(1,6)*kx2**2/beta0**2*k2pk1*(dx**3 - 2*sx*j2) + \
        Rational(1,2)*kx2**3/beta0**2*sx*j1 - Rational(1,2)*kx2/beta0**2*elrad*sx -\
        Rational(1,2)*kx2/beta0**2/gamma**2*dx
    t_iselin[0][2][2] = gk1*gk2*jd + Rational(1,2)*(gk2 + kx2*gk1)*dx
    t_iselin[0][2][3] = Rational(1,2)*gk2*js
    t_iselin[0][3][3] = gk2*jd - Rational(1,2)*kx2*dx
    
    t_iselin[1][0][0] = -Rational(1,6)*k2pk1*sx*(1 + 2*cx)
    t_iselin[1][0][1] = -Rational(1,6)*k2pk1*dx*(1 + 2*cx)
    t_iselin[1][1][1] = -Rational(1,3)*k2pk1*sx*dx - Rational(1,2)*kx2*sx
    t_iselin[1][0][5] = -Rational(1,12)*kx2/beta0*k2pk1*(3*cx*j1 + sx*dx) - \
        Rational(1,4)/beta0*gk1*(sx - elrad*cx)
    t_iselin[1][1][5] = -Rational(1,12)*kx2/beta0*k2pk1*(3*sx*j1 + dx**2) + \
        Rational(1,4)/beta0*gk1*elrad*sx
    t_iselin[1][5][5] = -Rational(1,6)*kx2**2/beta0**2*k2pk1*(sx*dx**2 - 2*cx*j2) - \
        Rational(1,2)*kx2/beta0**2*gk1*(cx*j1 - sx*dx) - Rational(1,2)*kx2/beta0**2/gamma**2*sx
    t_iselin[1][2][2] = gk1*gk2*js + Rational(1,2)*(gk2 + kx2*gk1)*sx
    t_iselin[1][2][3] = Rational(1,2)*gk2*jc
    t_iselin[1][3][3] = gk2*js - Rational(1,2)*kx2*sx
    
    t_iselin[2][0][2] = Rational(1,2)*gk2*(cy*jc - 2*gk1*sy*js) + Rational(1,2)*kx2*gk1*sx*sy
    t_iselin[2][0][3] = Rational(1,2)*gk2*(sy*jc - 2*cy*js) + Rational(1,2)*kx2*sx*cy
    t_iselin[2][1][2] = Rational(1,2)*gk2*(cy*js - 2*gk1*sy*jd) + Rational(1,2)*kx2*gk1*dx*sy
    t_iselin[2][1][3] = Rational(1,2)*gk2*(sy*js - 2*cy*jd) + Rational(1,2)*kx2*dx*cy
    t_iselin[2][2][5] = Rational(1,2)*kx2/beta0*gk2*(cy*jd - 2*gk1*sy*jf) + \
        Rational(1,2)*kx2**2/beta0*gk1*j1*sy - Rational(1,4)/beta0*gk1*elrad*sy
    t_iselin[2][3][5] = Rational(1,2)*kx2/beta0*gk2*(sy*jd - 2*cy*jf) + \
        Rational(1,2)*kx2**2/beta0*j1*cy - Rational(1,4)/beta0*(sy + elrad*cy)
        
    k2pk1b = gk2 + kx2*gk1
    
    t_iselin[3][0][2] = Rational(1,2)*gk1*gk2*(2*cy*js - sy*jc) + Rational(1,2)*k2pk1b*sx*cy
    t_iselin[3][0][3] = Rational(1,2)*gk2*(2*gk1*sy*js - cy*jc) + Rational(1,2)*k2pk1b*sx*sy
    t_iselin[3][1][2] = Rational(1,2)*gk1*gk2*(2*cy*jd - sy*js) + Rational(1,2)*k2pk1b*dx*cy
    t_iselin[3][1][3] = Rational(1,2)*gk2*(2*gk1*sy*jd - cy*js) + Rational(1,2)*k2pk1b*dx*sy
    t_iselin[3][2][5] = Rational(1,2)*kx2/beta0*gk1*gk2*(2*cy*jf - sy*jd) + \
        Rational(1,2)*kx2/beta0*k2pk1b*j1*cy + Rational(1,4)/beta0*gk1*(sy - elrad*cy)
    t_iselin[3][3][5] = Rational(1,2)*kx2/beta0*gk2*(2*gk1*sy*jf - cy*jd) + \
        Rational(1,2)*kx2/beta0*k2pk1b*j1*sy - Rational(1,4)/beta0*gk1*elrad*sy
        
    t_iselin[4][0][0] = Rational(1,12)*kx2/beta0*k2pk1*(sx*dx + 3*j1) - \
        Rational(1,4)/beta0*gk1*(elrad - sx*cx)
    t_iselin[4][0][1] = Rational(1,12)*kx2/beta0*k2pk1*dx**2 + Rational(1,4)/beta0*gk1*sx**2
    t_iselin[4][1][1] = Rational(1,6)*kx2/beta0*k2pk1*j2 - Rational(1,2)/beta0*sx - \
        Rational(1,4)/beta0*gk1*(j1 - sx*dx)
    t_iselin[4][0][5] = Rational(1,12)*kx2**2/beta0**2*k2pk1*(3*dx*j1 - 4*j2) + \
        Rational(1,4)*kx2/beta0**2*gk1*j1*(1 + cx) + Rational(1,2)*kx2/beta0**2/gamma**2*sx
    t_iselin[4][1][5] = Rational(1,12)*kx2**2/beta0**2*k2pk1*(dx**3 - 2*sx*j2) + \
        Rational(1,4)*kx2/beta0**2*gk1*sx*j1 + Rational(1,2)*kx2/beta0**2/gamma**2*dx
    t_iselin[4][5][5] = Rational(1,6)*kx2**3/beta0**3*k2pk1*(3*j3 - 2*dx*j2) + \
        Rational(1,6)*kx2**2/beta0**3*gk1*(sx*dx**2 - j2*(1 + 2*cx)) + \
        Rational(3,2)/beta0**3/gamma**2*(kx2**2*j1 - elrad)
    t_iselin[4][2][2] = -kx2/beta0*gk1*gk2*jf - Rational(1,2)*kx2/beta0*k2pk1b*j1 + \
        Rational(1,4)/beta0*gk1*(elrad - cy*sy)
    t_iselin[4][2][3] = -Rational(1,2)*kx2/beta0*gk2*jd - Rational(1,4)/beta0*gk1*sy**2
    t_iselin[4][3][3] = -kx2/beta0*gk2*jf + Rational(1,2)*kx2**2/beta0*j1 - \
        Rational(1,4)/beta0*(elrad + cy*sy)
        
        
    # since we have symmetry for the last 2 indices, we have to add the remaining terms
    for i in range(6):
        for j in range(6):
            for k in range(j):
                # the 3rd index is always greater than the 2nd one in the above definitions, therefore
                t_iselin[i][j][k] = t_iselin[i][k][j]  
    
    # now we compute the corresponding 2nd order for the drift-kick-drift operation. For this we need also
    # the 2nd-order drift entries at zero. Fortunately they are not many:
    di_drift2 = [[[0 for i in range(6)] for j in range(6)] for k in range(6)]  # initialization
    # remember that our drifts need to be defined with respect to 1/2 of elrad
    di_drift2[4][1][1] = elrad2
    di_drift2[4][3][3] = elrad2
    
    di_drift2[0][5][1] = elrad2
    di_drift2[0][1][5] = di_drift2[0][5][1]
    
    di_drift2[2][5][3] = elrad2
    di_drift2[2][3][5] = di_drift2[2][5][3]
    
    di_drift2[4][5][5] = 3*elrad2/gamma**2
    
    # option 6 mode
    [kick_map, kick_map2] = concat2(m_iselin, t_iselin, di_drift, di_drift2)
                
    # option 8 = drift-kick-drift mode
    [kick_res, kick_res2] = concat2(m_iselin, t_iselin, di_drift, di_drift2)
    
    # now we have
    #
    # kick_map, kick_map2
    #
    # and, correspondingly,
    #
    # m_iselin, t_iselin
    
if option == 9:
    # old version with formulas computed by hand. Perhaps there is an error here; the machine-computed
    # formulas in option 10 converge to the thick result nicely after ~ 8 slices, but remember they introduce
    # errors due to non-symplecticity.
    
    toshow = 'second-order thick map'
    
    dgx = gg_expr.diff(x)
    dgz = gg_expr.diff(z)

    S0_expr = S0(x, px, z, pz, sig, psig)
    
    # zero'th- and first order
    kick3 = [\
            x + elrad*(1 + kx2*x + kz2*z)*px/S0_expr, \
            px + elrad*(kx2*S0_expr + dgx), \
            z + elrad*(1 + kx2*x + kz2*z)*pz/S0_expr, \
            pz + elrad*(kz2*S0_expr + dgz), \
            sig + elrad(1 - (1 + kx2*x + kz2*z)*(1 + beta0**2*psig)/S0_expr), \
            psig\
           ]
    
    # add second-order
    dgxx = dgx.diff(x)
    dgxz = dgx.diff(z)
    dgzz = dgz.diff(z)
    
    kick3[0] = kick3[0] - elrad**2*Rational(1,2)*(1 + kx2*x + kz2*z)*(-3*kx2 - 4*kx2*px**2/S0_expr**2 - \
                                                                   3*(1 + px**2/S0_expr**2)/S0_expr*dgx - \
                                                                   4*kz2*px*pz/S0_expr**2 - 3*px*pz*dgz/S0_expr**3)
    
    kick3[1] = kick3[1] + elrad**2*Rational(1,2)*(-3*kx2**2*px - 3*kx2*px*dgx/S0_expr - 3*kx2*kz2*pz - 3*kx2*pz*dgz/S0_expr \
                                                - (1 + kx2*x + kz2*z)*(px*dgxx + dgxz*pz)/S0_expr)
    
    kick3[2] = kick3[2] - elrad**2*Rational(1,2)*(1 + kx2*x + kz2*z)*(-3*kz2 - 4*kz2*pz**2/S0_expr**2 - \
                                                                   3*(1 + pz**2/S0_expr**2)/S0_expr*dgz - \
                                                                   4*kx2*px*pz/S0_expr**2 - 3*px*pz*dgx/S0_expr**3)
    
    kick3[3] = kick3[3] + elrad**2*Rational(1,2)*(-3*kz2**2*pz - 3*kz2*pz*dgz/S0_expr - 3*kz2*kx2*px - 3*kz2*px*dgx/S0_expr \
                                                - (1 + kx2*x + kz2*z)*(pz*dgzz + dgxz*px)/S0_expr)
    
    kick3[4] = kick3[4] - elrad**2*Rational(1,2)*(1 + kx2*x + kz2*z)*(-4*kx2*px*pz/S0_expr**2 - 3*px*pz*dgx/S0_expr**3 \
                                                                   -3*kz2 -4*kz2*pz**2/S0_expr**2 - \
                                                                    3*dgz*(1 + pz**2/S0_expr**2)/S0_expr)
    
    ripken_used = kick3    
    
    
if option == 10:
    
    ham = hamil(x, px, z, pz, sig, psig, bx_hor_plane, bz_hor_plane)
    [dq, dp] = get_order(ham, 2)
    
    ripken_used = [dq[0], dp[0], dq[1], dp[1], dq[2], dp[2]]
    
if option == 11:
    
    ham = hamil(x, px, z, pz, sig, psig, bx_hor_plane, bz_hor_plane)
    [dq, dp] = get_kd_order(ham, 1)
    
    ripken_used = [dq[0], dp[0], dq[1], dp[1], dq[2], dp[2]]
    
    
if nshow:
    pure_latex = False
    
    print 'Transformation rules obtained by using the Hamiltonian\n'
    
    ham = hamil(x, px, z, pz, sig, psig, bx_hor_plane, bz_hor_plane)
    
    mystr = '\mathcal{H} = ' + latex(ham)
    if pure_latex:
        print mystr
    else:
        display(Math(mystr))

    print '\n*********\nOPTION:', option, ' ' + toshow + '\n*********'
    
    for i in range(6):
        mystr = str(ripken(i+1)) + '^f = ' + latex(ripken_used[i])
        if pure_latex:
            print mystr
        else:
            display(Math(mystr))

In [ ]:
def trf(i):
    if option not in [6, 7, 8]:
        vec = ripken_used
    
        # the new coordinates are orbit, not Ripken notation. So we have to
        # use our inverse madx-converter
        result = inv_madxconverter(orbit(i)).subs(ripken(i), vec[i-1])
        result = madxconverter(result)
    
        # we also have to replace kx and kz by the expressions used in the MAD-X code:
        #result = result.subs(kx, dipr/elrad).subs(kz, dipi/elrad)
    else:
        result = 0*x
    
    return N(result)
# the madx-converter assumes that its input is given in Ripken notation.



# compute the first and 2nd derivatives ('re' and 'te'-entries in src/twiss.f90)

def re(i, j, reference_point):

    
    if option == 6:   # option = 6 is Iselin drift-kick-drift at reference point = 0
        result = kick_map[i - 1][j - 1]
    elif option == 7:
        result = m_iselin[i - 1, j - 1]
    elif option == 8:
        result = kick_res[i - 1][j - 1]
    else:
        # in sympy we encounter problems if we simply derive orbit i with respect
        # to orbit j, because eta was originally a function of psig and the
        # result can not be easily substituted. Therefore we go back to ripken notation
        # derive everything there and go back to orbit notation.
    
        result = madxconverter(inv_madxconverter(trf(i)).diff(ripken(j)))
    
        # now we have to take into account that ripken(j) and orbit(j) might differ
        # by a factor. (We do not implement a general chain rule, since it will take
        # too much effort for nothing; the relation is fixed anyways)
    
        result = madxconverter(ripken(j)).subs(orbit(j), result)
        
    return expand(at_reference_point(result, reference_point))


def te(i, j, k, reference_point):

    
    if option == 6:   # option = 6 is Iselin drift-kick-drift at reference point = 0
        result = kick_map2[i - 1][j - 1][k - 1]
    elif option == 7:
        result = t_iselin[i - 1][j - 1][k - 1]
    elif option == 8:
        result = kick_res2[i-1][j-1][k-1]
    else:
        #d2_fac = 0.5  # 1/2!, necessary for 2nd order terms
    
        # first derivative wrt. j as in re
        result = madxconverter(inv_madxconverter(trf(i)).diff(ripken(j)))
        result = madxconverter(ripken(j)).subs(orbit(j), result)
        
        # second derivative wrt. k
        result = madxconverter(inv_madxconverter(result).diff(ripken(k)))
        result = madxconverter(ripken(k)).subs(orbit(k), result)

        d2_fac = 0.5
    
        result = at_reference_point(result, reference_point)*d2_fac
    
    return expand(result)


# in the fortran code, x stands for orbit(1) etc... we need to replace the symbols for the output

def fortran_strsub(string):
    
    step0 = string.replace('etah(p_sigma)', 'deltap')\
    .replace('sigma', 'sig')\
    .replace('p_x', 'px')\
    .replace('p_z', 'pz')\
    .replace('\Delta s', 'elrad')\
    .replace('K_x', '(dipr/elrad)')\
    .replace('K_z', '(dipi/elrad)')
    
    step1 = step0.replace('o1', 'x').replace('o2', 'px0').replace('o3', 'y')\
    .replace('o4', 'py0').replace('o5', 'orb50').replace('o6', 'orb60')\
    .replace('beta0', 'beta')\
    .replace('g', 'gstr')
            
    step2 = step1\
    .replace('etahat(orb60/beta)', 'deltap')\
    .replace('h_x', 'hfun')\
    .replace('h_z', 'hfun2')
    
    step3 = step2

    return step3

In [ ]:
# now we are able to produce the output
showzeros = False
indent = ''   # '\n'
show0, show1, show2 = False, True, True

try_simplify = True
print 'option =', option

if show0:
    print ''
    print '! orbit transformation:'
    
    for i in range(6):
        oi = trf(i+1).subs(kx, dipr/elrad).subs(kz, dipi/elrad)
        if try_simplify:
            oi = simplify(oi)
        if (oi != 0 or showzeros):
            m = '       orbit(' + str(i+1) + ') = ' + fortran_strsub(str(oi))
            print indent + f_wrapper(m)

if show1:
    print ''
    print '! first derivatives (at zero):'

    for i in range(6):
        for j in range(6):
            reij = re(i+1, j+1, [0]*6).subs(kx, dipr/elrad).subs(kz, dipi/elrad)
            if (reij != 0 or showzeros):
                m = '      re(' + str(i+1) + ', ' + str(j+1) + ') = ' + fortran_strsub(str(reij))
                print indent + f_wrapper(m)

if show2:
    print ''
    print '! 2nd derivatives (at zero):'
        
    for i in range(6):
        for j in range(6):
            for k in range(6):
                teijk = te(i+1, j+1, k+1, [0]*6).subs(kx, dipr/elrad).subs(kz, dipi/elrad)
                if (teijk != 0 or showzeros):
                    m = '        te(' + str(i+1) + ', ' + str(j+1) + ', ' + str(k+1) + ') = '\
                    + fortran_strsub(str(teijk))
                    print indent + f_wrapper(m)

In [ ]:
# Symplecticity check (at reference point rp):
rp = [0, 0, 0, 0, 0, 0]
re_expl = [[re(i+1, j+1, rp) for j in range(6)] for i in range(6)]   # re_expl[i][j] = re(i+1, j+1)

In [ ]:
J = Matrix(re_expl)
sym1 = Matrix([\
               [ 0,  1,  0,  0,  0,  0],\
               [-1,  0,  0,  0,  0,  0],\
               [ 0,  0,  0,  1,  0,  0],\
               [ 0,  0, -1,  0,  0,  0],\
               [ 0,  0,  0,  0,  0,  1],\
               [ 0,  0,  0,  0, -1,  0]\
              ])
matrix_to_check = J*sym1*J.transpose() - sym1

In [ ]:
display(Math(latex(simplify(J))))

In [ ]:
display(Math(latex(simplify(matrix_to_check))))

In [ ]:
def expl(expr, vals):
    result = expr
    if type(result) not in [int]:
        result = expr\
        .subs(kx, (dipr/elrad))\
        .subs(kz, (dipi/elrad))\
        .subs(dipr, vals[0])\
        .subs(dipi, vals[1])\
        .subs(g, vals[2]).subs(g2, vals[2])\
        .subs(beta0, vals[3])\
        .subs(elrad, vals[4])\
        .subs(sg, vals[5])
    
    return result

In [ ]:
nslices = 64 # Iselins thick formulas agree for 1 slice with the values of the MAD-X code


# These values are given according to the cfm.madx file (usually they do not require to be changed)
# check the MAD-X output, produced by cfm.madx, to obtain these values.
v_dipr = 0.39269908169872414
v_dipi = 0.0
v_g = -3.6443148688046642E-002   # quadrupole field strength
v_beta = 0.99995598518451123
v_elrad = 3.92
#v_sg = -4.0e-2     # sextupole field strength
v_sg = 0
vals = [v_dipr/nslices, v_dipi/nslices, v_g, v_beta, v_elrad/nslices, v_sg]

v_kx = v_dipr/v_elrad

print 'number of slices: ' + str(nslices)
print 'substituting ' + str(vals)

In [ ]:
option = 10

In [ ]:
rp = [0, 0, 0, 0, 0, 0]
re_expl = [[expl(re(i+1, j+1, rp), vals) for j in range(6)] for i in range(6)]   # re_expl[i][j] = re(i+1, j+1)
te_expl = [[[expl(te(i+1, j+1, k+1, rp), vals) for k in range(6)] for j in range(6)] for i in range(6)]  # te_expl[i][j][k] = te(i+1, j+1, k+1)

In [ ]:
def iten(n):
    # this function concatenates the first- and second derivatives of the form
    # h \circ g, where in every step h is known and g is unknown. I.e.
    #
    # step1: h \circ g
    # step2: h \circ (h \circ g)
    # step3: h \circ (h \circ (h \circ g))
    # etc. (\circ is associative)
    #
    # The first- and second derivatives of h correspond to re_expl and te_expl respectively.
    # The function was checked against the MAD-X output in order to verify the operations for
    # variable number of slices. So far all computed values agree very well.
    
    # iten is for example suitable for option = 7
    
    r_res = copy.deepcopy(re_expl)
    t_res = copy.deepcopy(te_expl)
    space = ' '*8
    print ''
    print 'slice:'
    for s in range(n):
        print  str(s + 1) +  \
        space + 'R11 = ' + str(r_res[0][0]) + \
        space + 'T111 = ' + str(t_res[0][0][0]) + \
        space + 'T511 = ' + str(t_res[4][0][0])   # t261 good test
        
        [r_res, t_res] = concat2(re_expl, te_expl, r_res, t_res)
        
        
def iten2(n):
    # same as iten, but now alternating drift- and kicks. Sometimes MAD-X makes 2 drifts in a row, this is not
    # implemented here. Note that the di_drift emerged from the inverse of a drift, so it must be inverted again.
    # The same goes for the 2nd derivative, where we have to replace elrad with -elrad (see my notes).
    
    # iten2 is suitable for option == 6. Note that it has to slice 2* the given number.
    re_expl_drift = copy.deepcopy([[expl(di_drift.inv()[i, j], vals) for j in range(6)] for i in range(6)])
    
    #print display(Math(latex(di_drift.inv())))
    
    # We need also
    # the 2nd-order drift entries at zero.
    d_drift2 = [[[0 for i in range(6)] for j in range(6)] for k in range(6)]  # initialization
    #zer = [[[0 for i in range(6)] for j in range(6)] for k in range(6)]  # initialization
    zer = copy.deepcopy(d_drift2)  # only with this operation, we will have a fresh clone.
    
    # remember that our drifts need to be defined with respect to 1/2 of elrad
    d_drift2[4][1][1] = -elrad2
    d_drift2[4][3][3] = -elrad2
    
    d_drift2[0][5][1] = -elrad2
    d_drift2[0][1][5] = d_drift2[0][5][1]
    
    d_drift2[2][5][3] = -elrad2
    d_drift2[2][3][5] = d_drift2[2][5][3]
    
    d_drift2[4][5][5] = -3*elrad2/gamma**2
    
    # print zer comment this in and use various copy commands to see that only deepcopy works
        
    
    te_expl_drift = copy.deepcopy([[[expl(d_drift2[i][j][k], vals) for k in range(6)] \
                                    for j in range(6)] for i in range(6)])
    
    r_res = copy.deepcopy(re_expl_drift)
    t_res = copy.deepcopy(te_expl_drift)    
    
    space = ' '*8
    print ''
    print 'slice:'
    for s in range(2*n):
        
        if s%2 == 1:
            my_r = re_expl_drift
            my_t = te_expl_drift
            #my_t = zer   # test
            action = ' drift'
        else:
            my_r = re_expl
            my_t = te_expl
            #my_t = zer
            action = ' kick '
        
        [r_res, t_res] = concat2(my_r, my_t, r_res, t_res)
        
        if s != 2*n - 1:  # dont print the values after the last kick
            print  str(s + 1) + action + \
            space + 'R11 = ' + str(r_res[0][0]) + \
            space + 'T111 = ' + str(t_res[0][0][0]) + \
            space + 'T511 = ' + str(t_res[4][0][0])

In [ ]:
iten2(nslices)

In [ ]:
iten2(nslices)

In [ ]:
iten2(nslices)

In [ ]:
iten(nslices)

In [ ]:
iten2(nslices)

In [ ]:
# thick lens comparison to Iselin (1985) = thick result in MAD-X

hise = v_dipr/v_elrad   # hise: curvature term h in his paper

kax = cmath.sqrt(hise**2 + v_g)
kay = cmath.sqrt(-v_g)

cx = cmath.cosh(1j*kax*v_elrad)
cy = cmath.cosh(1j*kay*v_elrad)

sx = cmath.sinh(1j*kax*v_elrad)/(1j*kax)
sy = cmath.sinh(1j*kay*v_elrad)/(1j*kay)

dx = (1.0 - cx)/kax**2

j1 = (v_elrad - sx)/kax**2

v_gamma = 1/math.sqrt(1 - v_beta**2)

m_iselin = Matrix([
        [cx, sx, 0, 0, 0, hise/v_beta*dx],
        [-kax**2*sx, cx, 0, 0, 0, hise/v_beta*sx],
        [0, 0, cy, sy, 0, 0],
        [0, 0, -kay**2*sy, cy, 0, 0],
        [-hise/v_beta*sx, -hise/v_beta*dx, 0, 0, 1, -(hise/v_beta)**2*j1 + v_elrad/(v_beta*v_gamma)**2],
        [0, 0, 0, 0, 0, 1]
    ])

In [ ]:
Math(latex(m_iselin))

In [ ]:
J = m_iselin
sym1 = Matrix([\
               [ 0,  1,  0,  0,  0,  0],\
               [-1,  0,  0,  0,  0,  0],\
               [ 0,  0,  0,  1,  0,  0],\
               [ 0,  0, -1,  0,  0,  0],\
               [ 0,  0,  0,  0,  0,  1],\
               [ 0,  0,  0,  0, -1,  0]\
              ])
matrix_to_check = J*sym1*J.transpose() - sym1

Math(latex(simplify(matrix_to_check)))
# OLD CODE! def ren(n, i, j): # returns re(i, j) of n-times concatenation if n == 1: sum1 = re_expl[i-1][j-1] if n > 1: sum1 = 0 for k in range(6): sum1 = sum1 + re_expl[i-1][k]*ren(n-1, k+1, j) return sum1 def ten(n, i, k, l): # returns te(i, j, k) of n-times concatenation if n == 1: sum2 = te_expl[i-1][k-1][l-1] if n > 1: sum1 = 0 for j in range(6): sum1 = sum1 + re_expl[i-1][j]*ten(n-1, j+1, k, l) sum2 = sum1 for a in range(6): for b in range(6): sum2 = sum2 + te_expl[i-1][a][b]*ren(n-1, a+1, k)*ren(n-1, b+1, l) return sum2|
def lamb(n): return (kx*(1 - 2*n) - I*kz*(1 - 2*n))*Rational(1,4*(n + 1)) def psi0(m, k, j): prod1 = 1 for w in range(j): # case j = 0 leads to range(0) = [] and this loop is not executed prod1 = prod1*lamb(m + w) prod2 = 1 for u in range(k - j): prod2 = prod2*conjugate(lamb(u)) C = functions.combinatorial.factorials.binomial return C(k, j)*prod1*prod2 def f_multipole(x, px, z, pz, sig, psig): # the initial conditions by which we determine the gauge of the potential bx_hor_plane = [kz] bz_hor_plane = [kx, g] # get gamma-conditions vector, n0 and N (see theory) gamma = bz_hor_plane j = 0 n0 = 0 for ele in bx_hor_plane: try: gamma[j] = gamma[j] + I*ele except IndexError: gamma.append(I*ele) if (gamma[j] == 0) and (n0 == j): n0 = n0 + 1 j = j + 1 N = len(gamma) - 1 # gamma : complex initial condition vector # n0, N : lowest and maximal multipole order (+1) which are not zero # the following values are initialized and must not be changed! # mo = multipole order: current running index of ambm # f = function F (see theory) rp = x + I*z rm = x - I*z f = 0 An_all = [] # list of all An's, starting with A_(n0+1), i.e. An_all = [A_(n0+1), ... A_(N+2)] for n in range(n0, N + 2): # n = n0, 1, ..., N + 1 # compute A_(n+1) and B_(n+1) coefficients (see theory) Bnp1 = 0 if (n0 == n): Bnp1 = gamma[n0] if (n0 < n) and (n <= N): Bnp1 = gamma[n] + kx*gamma[n - 1] if (n == N + 1): Bnp1 = kx*gamma[N] Bnp1 = -Rational(1,2*(n + 1))*conjugate(Bnp1) # the conjugation fixes a bug in a previous version for m in range(n0 + 1, n + 1): # m = n0 + 1, ..., n. This loop is not entered if n = n0. omega_p = 0 Anp1_m = An_all[m - n0 - 1] for j in range(n - m + 2): # j = 0, ..., n - m + 1 psi = Anp1_m*psi0(m, n - m + 1, j) psi_tilde = conjugate(Anp1_m*psi0(m, n - m + 1, n - m + 1 - j)) omega_p = omega_p + (n - m + 1 - j)*psi + (n + 1 - j)*psi_tilde Bnp1 = Bnp1 - Rational(1,n + 1)*omega_p Anp1 = conjugate(Bnp1) An_all.append(Anp1) # from here on, A_(n+1) and B_(n+1) are determined. # every multipole of order n+1 generates higher orders. Since we consider multipole orders of # up to N+2 (n goes up to N + 1 and thus the last coeffecients are A_(N+2), B_(N+2)), # the current multipole of order n+1 must at least *) be developed up to N - n + 1 times with respect # to the development order k (F = sum_k F_k, see theory). # *) If higher orders are of interest, add additional orders by setting n_ord to an integer greater than zero. # For n_ord \to \infty we obtain a Maxwell-conformal field, however the conditions in the midplane z = 0 # for these higher orders are not fixed by any gauge. Thus it doesnt make much sense to set a value # greater than zero. Instead, one should fix the gauge by adding zeros (or other values) to the # midplane-conditions bx_hor_plane and bz_hor_plane according to what field one requires. n_ord = 0 sumk = 0 for k in range(N - n + 2 + n_ord): # k = 0, ..., N - n + 1 + n_ord sumj = 0 for j in range(k + 1): # j = 0, ..., k sumj = sumj + psi0(n + 1, k, j)*rp**j*rm**(k - j) sumk = sumk + sumj f = f + 2*re(Anp1*rp**(n+1)*sumk) return expand(f)

In [ ]:
for i in range(6):
    for j in range(6):
        print expl(matrix_to_check[i, j].subs(o1, 0.3322), vals)

In [ ]: